package com.zero.manage.web.utils.dynamicsource;

import cn.hutool.core.util.ClassUtil;
import cn.hutool.core.util.StrUtil;
import com.baomidou.dynamic.datasource.DynamicRoutingDataSource;
import com.baomidou.dynamic.datasource.annotation.DS;
import com.baomidou.dynamic.datasource.creator.DruidDataSourceCreator;
import com.baomidou.dynamic.datasource.spring.boot.autoconfigure.DataSourceProperty;
import io.gitee.zerowsh.enums.ToolEnums;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.sql.DataSource;
import java.lang.reflect.Method;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.HashSet;
import java.util.Objects;
import java.util.Set;

@Component
@AllArgsConstructor
@Slf4j
public class DynamicSourceUtil {

    private final DruidDataSourceCreator druidDataSourceCreator;
    private final DynamicRoutingDataSource ds;

    /**
     * 初始化数据源
     */
    public void initSource() {
        Set<String> sourceSet = new HashSet<>();
        Set<Class<?>> classes = ClassUtil.scanPackage("com.zero.manage.web.components.service.impl");
        for (Class<?> aClass : classes) {
            //处理类上注解
            DS classAnnotation = aClass.getAnnotation(DS.class);
            if (!Objects.isNull(classAnnotation) && StrUtil.isNotBlank(classAnnotation.value())) {
                String value = classAnnotation.value();
                if (!value.contains("#")) {
                    sourceSet.add(value);
                }
            }
            //处理方法上注解
            Method[] methods = aClass.getMethods();
            for (Method method : methods) {
                DS methodAnnotation = method.getAnnotation(DS.class);
                if (!Objects.isNull(methodAnnotation) && StrUtil.isNotBlank(methodAnnotation.value())) {
                    String value = methodAnnotation.value();
                    if (!value.contains("#")) {
                        sourceSet.add(value);
                    }
                }
            }
        }
        for (String sourceId : sourceSet) {
            addSource(sourceId);
        }
    }

    /**
     * 添加数据源
     *
     * @param sourceId
     */
    public void addSource(String sourceId) {
        if (StrUtil.isNotBlank(sourceId)) {
            if (!ds.getCurrentDataSources().containsKey(sourceId)) {
                ResultSet rs = null;
                try (Connection conn = ds.getConnection();
                     PreparedStatement ps = conn.prepareStatement("select database_type,username,password,ip_port,other_params from t_database_info where id=?")) {
                    ps.setString(1, sourceId);
                    rs = ps.executeQuery();
                    //获取查询结果集
                    while (rs.next()) {
                        DataSourceProperty dataSourceProperty = new DataSourceProperty();
                        dataSourceProperty.setUsername(rs.getString("username"));
                        dataSourceProperty.setPassword(rs.getString("password"));
                        dataSourceProperty.setUrl(ToolEnums.DataBaseEnums.getUrlPrefixByName(rs.getString("database_type"))
                                + rs.getString("ip_port") + rs.getString("other_params"));
                        DataSource dataSource = druidDataSourceCreator.createDataSource(dataSourceProperty);
                        ds.addDataSource(sourceId, dataSource);
                    }
                } catch (Exception e) {
                    log.error("动态添加数据源失败", e);
                } finally {
                    if (!Objects.isNull(rs)) {
                        try {
                            rs.close();
                        } catch (SQLException e) {
                            log.error("关闭ResultSet失败", e);
                        }
                    }
                }
            }
        }
    }
}
